import cv2 as cv
import numpy as np


def watershed_demo(image):
    blurred = cv.pyrMeanShiftFiltering(image,10,100)# sobel_demo(image)
    # gray,binary
    gray = cv.cvtColor(blurred,cv.COLOR_BGR2GRAY)
    ret,binary = cv.threshold(gray,0,255,cv.THRESH_BINARY|cv.THRESH_OTSU)
    cv.imshow('binary',binary)

    # morphology
    kernel = cv.getStructuringElement(cv.MORPH_RECT,(3,3))
    mb = cv.morphologyEx(binary,cv.MORPH_OPEN,kernel,iterations=2)
    cv.imshow('open-binary',mb)

    # distance transform
    dist = cv.distanceTransform(mb,cv.DIST_L2,3)
    dist_out = cv.normalize(dist,0,1.0,cv.NORM_MINMAX)
    cv.imshow('distance',dist_out*50)

    ret, surface = cv.threshold(dist,dist.max()*0.6,255,cv.THRESH_BINARY)
    cv.imshow('surface-bin',surface)

    sure_bg = cv.dilate(mb,kernel,iterations=3)
    surface_fg =np.uint8(surface)
    unknown = cv.subtract(sure_bg,surface_fg)

    ret,markers = cv.connectedComponents(surface_fg)
    # watershed
    markers = markers + 1
    markers[unknown==255] = 0
    markers = cv.watershed(image,markers)
    image[markers == -1] = [0,0,255]
    cv.imshow('watershed',image)


def sobel_demo(image):
    grad_x = cv.Sobel(image,cv.CV_32F,1,0)
    grad_y = cv.Sobel(image,cv.CV_32F,0,1)
    gradx = cv.convertScaleAbs(grad_x)
    grady = cv.convertScaleAbs(grad_y)
    added = cv.addWeighted(gradx,0.5,grady,0.5,0)
    return added


src = cv.imread('coins.jpg',1)
cv.namedWindow('demo',cv.WINDOW_AUTOSIZE)
cv.imshow('demo',src)
watershed_demo(src)
cv.waitKey(0)
cv.destroyAllWindows()